import networkx as nx
import pandas as pd
import matplotlib.pyplot as plt
import community as community_louvain
from matplotlib.font_manager import FontProperties

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号'-'显示为方块的问题

# 读取数据
entities_df = pd.read_csv('./data/entity.csv')
relationships_df = pd.read_csv('./data/relationship.csv')

# 创建图
G = nx.Graph()

# 添加节点
for index, row in entities_df.iterrows():
    G.add_node(row['id'])

# 添加边
for index, row in relationships_df.iterrows():
    G.add_edge(row['Source'], row['Target'], weight=row['Weight'])

# 使用 Fast Unfolding 算法进行社区发现
partition = community_louvain.best_partition(G)


# 选择颜色映射
cmap_name = 'autumn'  # 可以替换为 'plasma', 'inferno', 'magma', 'cividis' 等
cmap = plt.cm.get_cmap(cmap_name, max(partition.values()) + 1)

# 为每个社区分配颜色
community_colors = {com: cmap(float(com) / max(partition.values())) for com in set(partition.values())}


# 使用networkx的spring_layout布局算法设置节点位置
pos = nx.spring_layout(G, k=0.15, iterations=1000)  # 调整布局参数以获得更好的视觉效果

# 绘制图谱
plt.figure(figsize=(12, 12))
for com in set(partition.values()):
    nodes_in_com = [nodes for nodes in partition.keys() if partition[nodes] == com]
    nx.draw_networkx_nodes(G, pos, nodelist=nodes_in_com, node_color=community_colors[com], node_size=300)  # 调整节点大小
nx.draw_networkx_edges(G, pos, edge_color="black", alpha=0.05, width=0.5)  # 调整边的宽度和透明度
nx.draw_networkx_labels(G, pos, font_size=14, font_family="sans-serif", font_color="black")  # 调整标签字体大小和颜色

# 绘制社区的轮廓
for com in set(partition.values()):
    nodes_in_com = [nodes for nodes in partition.keys() if partition[nodes] == com]
    internal_edges = [(nodes_in_com[i], nodes_in_com[j]) for i in range(len(nodes_in_com)) for j in range(i+1, len(nodes_in_com)) if G.has_edge(nodes_in_com[i], nodes_in_com[j])]
    if internal_edges:
        internal_subgraph = G.subgraph(nodes_in_com)
        internal_pos = {node: pos[node] for node in internal_subgraph.nodes()}
        nx.draw_networkx_edges(internal_subgraph, internal_pos, edgelist=internal_edges, edge_color=community_colors[com], width=1, alpha=0.3)  # 调整社区内边的宽度和透明度

plt.title('知识图谱与社区', fontdict={'fontsize': 20, 'fontweight': 'medium', 'color': 'navy'})  # 调整标题字体大小、粗细和颜色
plt.axis('off')
plt.legend()
plt.savefig('./data2/知识图谱与社区.png', dpi=300)  # 保存高分辨率图像
plt.show()